/*
 * message_passing.cc
 *
 *  Created on: Sep 12, 2011
 *      Author: tqlong
 */

#include <boost/foreach.hpp>
#include "message_passing.h"

MessagePassingGM::MessagePassingGM(const SimpleGM& gm)
{
  cardinality_ = gm.cardinality_;
  node_potentials_ = gm.node_potentials_;
  edge_potentials_ = gm.edge_potentials_;
  adj_edges_ = gm.adj_edges_;
  node_marginals_ = gm.node_marginals_;
  edge_marginals_ = gm.edge_marginals_;
  old2new_ = gm.old2new_;
  new2old_ = gm.new2old_;
}

void MessagePassingGM::message_passing_marginals()
{
  init_messages();

  // pass messages until convergence
  for (unsigned int iter = 0; iter < messages_.size(); iter++)
  {
    message_type tmp_msg = messages_;
    create_passing_schedule();
    BOOST_FOREACH(const Edge& e, schedule_)
    {
      update_message(e, tmp_msg);  // use current messages_ to compute new message in tmp_msg
    }
    double d = message_diffence(messages_, tmp_msg);
    std::cout << "iter = " << iter << std::endl;
    std::cout << "Changes = " << d << std::endl;

    messages_ = tmp_msg;           // update new messages_
    print_messages();

    if (d < 1e-5)
    {
      std::cout << "Changes too small !!!" << std::endl;
      break;
    }
  }

  // compute marginals
  init_marginals();
  message_to_node_marginals();   // node marginals
  message_to_edge_marginals();   // edge marginals
}

MessagePassingGM::assignment_type MessagePassingGM::message_passing_map()
{
  init_messages();
  // pass messages until convergence
  for (unsigned int iter = 0; iter < messages_.size(); iter++)
  {
    message_type tmp_msg = messages_;
    create_passing_schedule();
    BOOST_FOREACH(const Edge& e, schedule_)
    {
      update_max_message(e, tmp_msg);  // use current messages_ to compute new message in tmp_msg
    }
    double d = message_diffence(messages_, tmp_msg);
    std::cout << "iter = " << iter << std::endl;
    std::cout << "Changes = " << d << std::endl;

    messages_ = tmp_msg;           // update new messages_
    print_messages();

    if (d < 1e-5)
    {
      std::cout << "Changes too small !!!" << std::endl;
      break;
    }
  }

  message_to_max_marginals();
  assignment_type a;
  BOOST_FOREACH(const_pair_node_t p, max_marginals_)
  {
    Node i = p.first;
    const arma::vec& v = p.second; // the max marginal
    int val = 0;
    for (int vi = 1; vi < cardinality_[i]; vi++)
      if (v[vi] > v[val]) val = vi;
    a[i] = val;
  }
  return a;
}

void MessagePassingGM::print_messages()
{
  BOOST_FOREACH(const_pair_message_t p, messages_)
  {
    std::cout << "m(" << p.first.i_ << "," << p.first.j_ << ") ="
              << std::endl << p.second;
  }
}

// Protected members
void MessagePassingGM::init_messages()
{
  BOOST_FOREACH(const_pair_edge_t p, edge_potentials_)
  {
    Edge e = p.first;
    create_message(e.i_, e.j_);
    if (e.i_ != e.j_)
      create_message(e.j_, e.i_);
  }
}

void MessagePassingGM::create_message(int i, int j)
{
  int card = cardinality_[j];
  arma::vec v(card);
  v.fill(log(1));
  messages_[Edge(i, j, true)] = v;
}

// Compute message from i->j given message k->i \forall k\in N(i)\j
// use adj_edges to
void MessagePassingGM::update_message(const Edge& e, message_type& msg)
{
  int i = e.i_, j = e.j_;
  arma::vec& v = msg.find(e)->second;  // the message, need to fill this vector with new values
  v.fill(log(0));                      // start the accumulation
  // vector version
  arma::vec tmp = node_potentials_[i];
  BOOST_FOREACH(const Edge& e, adj_edges_.find(i)->second) // on all neighbor edge
  {
    int k = e.other_end(i);  // the neighbor of i on edge e
    if (k == j) continue;    // only consider k \in N(i) \ {j}
    Edge e_ki(k, i, true);
    arma::vec& v_ki = messages_.find(e_ki)->second; // old message from k -> i
    tmp += v_ki;
  }

  for (int vj = 0; vj < cardinality_[j]; vj++)
  {
    arma::mat& pot_ij = edge_potentials_[Edge(i, j)];
    arma::vec s;
    if (i < j)
      s = tmp + pot_ij.col(vj);
    else
      s = tmp + arma::trans(pot_ij.row(vj));
    v.at(vj) = logsum(s);
  }

//  //std::cout << i << " --> " << j << std::endl;
//  for (int vj = 0; vj < cardinality_[j]; vj++)
//  {
//    for (int vi = 0; vi < cardinality_[i]; vi++)
//    {
//      double tmp = node_potentials_[i].at(vi) + edge_potentials(i, j, vi, vj);
//      BOOST_FOREACH(const Edge& e, adj_edges_.find(i)->second) // on all neighbor edge
//      {
//        int k = e.other_end(i);  // the neighbor of i on edge e
//        if (k == j) continue;    // only consider k \in N(i) \ {j}
//        Edge e_ki(k, i, true);
//        arma::vec& v_ki = messages_.find(e_ki)->second; // old message from k -> i
//        tmp += v_ki.at(vi);
//      }
//      //std::cout << "    " << vi << "," << vj << " = " << tmp << std::endl;
//      v.at(vj) = logadd(v.at(vj), tmp);
//    }
//  }
}

void MessagePassingGM::update_max_message(const Edge& e, message_type& msg)
{
  int i = e.i_, j = e.j_;
  arma::vec& v = msg.find(e)->second;  // the message, need to fill this vector with new values
  v.fill(log(0));                      // start the accumulation
  // vector version
  arma::vec tmp = node_potentials_[i];
  BOOST_FOREACH(const Edge& e, adj_edges_.find(i)->second) // on all neighbor edge
  {
    int k = e.other_end(i);  // the neighbor of i on edge e
    if (k == j) continue;    // only consider k \in N(i) \ {j}
    Edge e_ki(k, i, true);
    arma::vec& v_ki = messages_.find(e_ki)->second; // old message from k -> i
    tmp += v_ki;
  }

  for (int vj = 0; vj < cardinality_[j]; vj++)
  {
    arma::mat& pot_ij = edge_potentials_[Edge(i, j)];
    arma::vec s;
    if (i < j)
      s = tmp + pot_ij.col(vj);
    else
      s = tmp + arma::trans(pot_ij.row(vj));
    v.at(vj) = arma::max(s);
  }
}

// Create a passing schedule
// for now, it is all message in order inside messages_
void MessagePassingGM::create_passing_schedule()
{
  schedule_.clear();
  BOOST_FOREACH(const_pair_message_t p, messages_)
  {
    schedule_.push_back(p.first);
  }
}

double MessagePassingGM::message_diffence(const message_type& m1, const message_type& m2)
{
  double s = 0;
  BOOST_FOREACH(const_pair_message_t p, m1)
  {
    const arma::vec& v1 = p.second;
    const arma::vec& v2 = m2.find(p.first)->second;
    s += arma::norm(v1 - v2, 2);
  }
  return s;
}

void MessagePassingGM::message_to_node_marginals()
{
  BOOST_FOREACH(pair_node_t p, node_marginals_)
  {
    Node i = p.first;
    arma::vec& v = p.second; // the marginals (need to fill)
    arma::vec& p_i = node_potentials_.find(i)->second; // the node potential
    v = p_i;
    BOOST_FOREACH(const Edge& e, adj_edges_[i])
    {
      int j = e.other_end(i);
      arma::vec& m_ji = messages_.find(Edge(j,i,true))->second;
      v += m_ji;
    }
  }
}

void MessagePassingGM::message_to_edge_marginals()
{
  BOOST_FOREACH(pair_edge_t p, edge_marginals_)
  {
    Edge e = p.first;
    int i = e.i_, j = e.j_;
    arma::mat& v = p.second; // the marginals (need to fill)
    v = edge_potentials_[e];

    arma::vec tmp_i(cardinality_[i]), tmp_j(cardinality_[j]);
    tmp_i.fill(log(1));
    tmp_j.fill(log(1));

    BOOST_FOREACH(const Edge& ei, adj_edges_[i])
    {
      int k = ei.other_end(i);
      if (k == j) continue;
      arma::vec& m_ki = messages_.find(Edge(k,i,true))->second;
      tmp_i += m_ki;
    }

    BOOST_FOREACH(const Edge& ej, adj_edges_[j])
    {
      int k = ej.other_end(j);
      if (k == i) continue;
      arma::vec& m_kj = messages_.find(Edge(k,j,true))->second;
      tmp_j += m_kj;
    }

    for (int vi = 0; vi < cardinality_[i]; vi++)
    {
      for (int vj = 0; vj < cardinality_[j]; vj++)
      {
        v.at(vi, vj) += tmp_i.at(vi) + node_potentials_[i].at(vi) +
                        tmp_j.at(vj) + node_potentials_[j].at(vj);
      }
    }
  }
}

void MessagePassingGM::message_to_max_marginals()
{
  max_marginals_ = node_potentials_;
  BOOST_FOREACH(pair_node_t p, max_marginals_)
  {
    Node i = p.first;
    arma::vec& v = p.second; // the marginals (need to fill)
    arma::vec& p_i = node_potentials_[i]; // the node potential
    v = p_i;
    BOOST_FOREACH(const Edge& e, adj_edges_[i])
    {
      int j = e.other_end(i);
      arma::vec& m_ji = messages_[Edge(j,i,true)];
      v += m_ji;
    }
  }
}
